# ******************************************************************************
#  Copyright (c) 2023 Orbbec 3D Technology, Inc
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http:# www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
# ******************************************************************************
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from orbbec.pyorbbecsdk import *
import cv2
import numpy as np
from typing import List, Optional
from queue import Queue

MAX_DEVICES = 2  # 最大支持设备数量
MAX_QUEUE_SIZE = 5  # 队列最大长度
ESC_KEY = 27  # ESC 键的 ASCII 码


class InfraredCamera:
    def __init__(self):
        self.ctx = Context()
        self.device_list = self.ctx.query_devices()
        self.curr_device_cnt = self.device_list.get_count()
        self.pipelines: List[Pipeline] = []
        self.configs: List[Config] = []
        self.has_ir_sensor: List[bool] = [False for _ in range(MAX_DEVICES)]
        self.ir_frames_queue: List[Queue] = [Queue() for _ in range(MAX_DEVICES)]
        self.stop_rendering = False

    def initialize(self):
        """
        初始化设备
        """
        if self.curr_device_cnt == 0:
            print("No device connected")
            return False
        if self.curr_device_cnt > MAX_DEVICES:
            print("Too many devices connected")
            return False
        for i in range(self.curr_device_cnt):
            device = self.device_list.get_device_by_index(i)
            pipeline = Pipeline(device)
            config = Config()
            try:
                profile_list = pipeline.get_stream_profile_list(OBSensorType.IR_SENSOR)
                ir_profile: VideoStreamProfile = profile_list.get_default_video_stream_profile()
                config.enable_stream(ir_profile)
                self.has_ir_sensor[i] = True
            except OBError as e:
                print(e)
                self.has_ir_sensor[i] = False
            self.pipelines.append(pipeline)
            self.configs.append(config)
        return True

    def start(self):
        """
        启动设备流
        """
        for index, (pipeline, config) in enumerate(zip(self.pipelines, self.configs)):
            print("Starting device {}".format(index))
            pipeline.start(config, lambda frame_set, curr_index=index: self.on_new_frame_callback(frame_set, curr_index))

    def stop(self):
        """
        停止设备流
        """
        for pipeline in self.pipelines:
            pipeline.stop()

    def on_new_frame_callback(self, frames: FrameSet, index: int):
        """
        回调函数，处理新帧数据
        """
        ir_frame = frames.get_ir_frame()  # 获取红外帧
        if ir_frame is not None:
            if self.ir_frames_queue[index].qsize() >= MAX_QUEUE_SIZE:
                self.ir_frames_queue[index].get()  # 如果队列满了，移除最旧的帧
            self.ir_frames_queue[index].put(ir_frame)  # 将新帧加入队列

    def get_ir_data(self, device_index: int) -> Optional[np.ndarray]:
        """
        获取红外数据
        :param device_index: 设备索引
        :return: 红外数据 (numpy 数组)，如果无数据则返回 None
        """
        if device_index >= self.curr_device_cnt or not self.has_ir_sensor[device_index]:
            return None
        if self.ir_frames_queue[device_index].empty():
            return None
        ir_frame = self.ir_frames_queue[device_index].get()  # 从队列中获取红外帧
        width = ir_frame.get_width()
        height = ir_frame.get_height()
        ir_data = np.asanyarray(ir_frame.get_data())
        ir_format = ir_frame.get_format()

        if ir_format == OBFormat.Y8:
            ir_data = np.resize(ir_data, (height, width, 1))
            data_type = np.uint8
            image_dtype = cv2.CV_8UC1
            max_data = 255
        elif ir_format == OBFormat.MJPG:
            ir_data = cv2.imdecode(ir_data, cv2.IMREAD_UNCHANGED)
            data_type = np.uint8
            image_dtype = cv2.CV_8UC1
            max_data = 255
            if ir_data is None:
                print("Decode MJPEG failed")
                return None
            ir_data = np.resize(ir_data, (height, width, 1))
        else:
            ir_data = np.frombuffer(ir_data, dtype=np.uint16)
            data_type = np.uint16
            image_dtype = cv2.CV_16UC1
            max_data = 65535
            ir_data = np.resize(ir_data, (height, width, 1))

        # 归一化数据
        cv2.normalize(ir_data, ir_data, 0, max_data, cv2.NORM_MINMAX, dtype=image_dtype)
        ir_data = ir_data.astype(data_type)
        return ir_data

    def display(self):
        """
        显示红外图像
        """
        while not self.stop_rendering:
            for i in range(self.curr_device_cnt):
                ir_data = self.get_ir_data(i)
                if ir_data is None:
                    continue
                ir_image = cv2.cvtColor(ir_data, cv2.COLOR_GRAY2RGB)  # 转换为 RGB 图像
                cv2.imshow("Infrared Device {}".format(i), ir_image)
            key = cv2.waitKey(1)
            if key == ord('q') or key == ESC_KEY:
                self.stop_rendering = True
                break


def main():
    """
    主函数
    """
    infrared_camera = InfraredCamera()
    if not infrared_camera.initialize():
        return
    infrared_camera.start()
    try:
        infrared_camera.display()  # 显示红外图像
    except KeyboardInterrupt:
        infrared_camera.stop_rendering = True
    finally:
        infrared_camera.stop()  # 停止设备流


if __name__ == "__main__":
    infrared_camera = InfraredCamera()
    if infrared_camera.initialize():
        infrared_camera.start()
        try:
            while True:
                ir_data1 = infrared_camera.get_ir_data(device_index=0)
                ir_data2 = infrared_camera.get_ir_data(device_index=1)
                if ir_data1 is None or ir_data2 is None:
                    continue
                ir_image1 = cv2.cvtColor(ir_data1, cv2.COLOR_GRAY2RGB)  # 转换为 RGB 图像
                ir_image2 = cv2.cvtColor(ir_data2, cv2.COLOR_GRAY2RGB)  # 转换为 RGB 图像
                cv2.imshow("Infrared Device1 ", ir_image1)
                cv2.imshow("Infrared Device2 ", ir_image2)
                cv2.waitKey(1)
        except KeyboardInterrupt:
            pass
        finally:
            infrared_camera.stop()